#include "aes.h"

// sets the key to use.
// the actual internal key will be looped if not long enough.
// that is, a password of "abcd" will actually create a key
// like "abcdabcdabcdabcd..."
void AES::setkey(std::string password)
{
    for(unsigned int i = 0; i < sizeof(key); i++)
        key[i] = password[i % password.length()];
}

// encrypts string 'in' into buffer 'out'
// inlen should be the length of the in buffer, or length of the string in that buffer
// outlen should be the length of the out buffer
// returns the full length of the encrypted string
//
// if you want to encrypt a null terminated string, make sure you use
// strlen(mystring)+1 for inlen (to include the null terminator)
unsigned int AES::encrypt(const char *in, char *out, unsigned int inlen, unsigned int outlen)
{
    unsigned int nrounds = rijndaelSetupEncrypt(rk, key, 256);

    unsigned int i = 0;
    unsigned int buffpos = 0;
    while( i < inlen )
    {
        unsigned char plaintext[16];
        unsigned char ciphertext[16];

        unsigned int j;
        for(j = 0; j < sizeof(plaintext); j++)
        {
            int c;
            if( i < inlen ) {
                c = in[i];
                i++; }
            else
                c = 0;

            plaintext[j] = c;
        }

        rijndaelEncrypt(rk, nrounds, plaintext, ciphertext);
        for(j = 0; j < sizeof(ciphertext); j++)
        {
            if( buffpos < outlen )
            {
                out[buffpos] = ciphertext[j];
                buffpos++;
            }
            else
            { // break so we aren't running over the out buffer
                i = inlen;
                break;
            }
        }
    }

    return buffpos;
}

// decrypts string 'in' into buffer 'out'
// inlen should be the length of the in buffer, or length of the string in that buffer
// outlen should be the length of the out buffer
// returns the full length of the decrypted string
unsigned int AES::decrypt(const char *in, char *out, unsigned int inlen, unsigned int outlen)
{
    unsigned int nrounds = rijndaelSetupDecrypt(rk, key, 256);

    unsigned int i = 0;
    unsigned int buffpos = 0;

    while( i < inlen )
    {
        unsigned char plaintext[16];
        unsigned char ciphertext[16];

        unsigned int j;
        for(j = 0; j < sizeof(ciphertext); j++)
        {
            int c;
            if( i < inlen ) {
                c = in[i];
                i++; }
            else
                c = 0;

            ciphertext[j] = c;
        }

        rijndaelDecrypt(rk, nrounds, ciphertext, plaintext);
        for(j = 0; j < sizeof(plaintext); j++)
        {
            if( buffpos < outlen )
            {
                out[buffpos] = plaintext[j];
                buffpos++;
            }
            else
            { // break so we aren't running over the out buffer
                i = inlen;
                break;
            }
        }
    }

    return buffpos;
}
